#!/bin/bash

# 容器健康检查脚本
# 用于Docker容器健康检查

set -e

# 配置
HEALTH_CHECK_URL="http://localhost:8000/api/v1/health"
TIMEOUT=10
RETRY_COUNT=3
RETRY_INTERVAL=2

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color

# 日志函数
log_info() {
    echo -e "${GREEN}[INFO]${NC} $1"
}

log_warn() {
    echo -e "${YELLOW}[WARN]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# 检查命令是否存在
check_command() {
    if ! command -v $1 &> /dev/null; then
        log_error "Command '$1' not found"
        exit 1
    fi
}

# HTTP健康检查
http_health_check() {
    local attempt=1
    local success=false
    
    while [ $attempt -le $RETRY_COUNT ]; do
        log_info "Health check attempt $attempt/$RETRY_COUNT"
        
        if curl -f -s --max-time $TIMEOUT "$HEALTH_CHECK_URL" > /dev/null 2>&1; then
            log_info "Health check passed"
            success=true
            break
        else
            log_warn "Health check failed, retrying in ${RETRY_INTERVAL}s..."
            sleep $RETRY_INTERVAL
        fi
        
        attempt=$((attempt + 1))
    done
    
    if [ "$success" = false ]; then
        log_error "Health check failed after $RETRY_COUNT attempts"
        return 1
    fi
    
    return 0
}

# 数据库连接检查
database_health_check() {
    log_info "Checking database connection..."
    
    # 这里应该检查数据库连接
    # 由于是SQLite，检查文件是否存在且可访问
    if [ -f "/app/data/app.db" ]; then
        log_info "Database file exists"
        
        # 尝试连接数据库
        python3 -c "
import sqlite3
import sys
try:
    conn = sqlite3.connect('/app/data/app.db', timeout=5)
    cursor = conn.cursor()
    cursor.execute('SELECT 1')
    conn.close()
    print('Database connection successful')
except Exception as e:
    print(f'Database connection failed: {e}', file=sys.stderr)
    sys.exit(1)
" 2>/dev/null
        
        if [ $? -eq 0 ]; then
            log_info "Database connection check passed"
            return 0
        else
            log_error "Database connection check failed"
            return 1
        fi
    else
        log_warn "Database file not found, attempting to create..."
        python3 -c "
import sqlite3
import sys
try:
    conn = sqlite3.connect('/app/data/app.db')
    conn.close()
    print('Database created successfully')
except Exception as e:
    print(f'Database creation failed: {e}', file=sys.stderr)
    sys.exit(1)
" 2>/dev/null
        
        if [ $? -eq 0 ]; then
            log_info "Database created successfully"
            return 0
        else
            log_error "Database creation failed"
            return 1
        fi
    fi
}

# Redis连接检查
redis_health_check() {
    log_info "Checking Redis connection..."
    
    python3 -c "
import redis
import sys
import os

redis_url = os.getenv('REDIS_URL', 'redis://localhost:6379/0')
try:
    r = redis.from_url(redis_url)
    r.ping()
    print('Redis connection successful')
except Exception as e:
    print(f'Redis connection failed: {e}', file=sys.stderr)
    sys.exit(1)
" 2>/dev/null
    
    if [ $? -eq 0 ]; then
        log_info "Redis connection check passed"
        return 0
    else
        log_error "Redis connection check failed"
        return 1
    fi
}

# 内存使用检查
memory_health_check() {
    log_info "Checking memory usage..."
    
    # 获取内存使用情况
    memory_info=$(free -m | grep Mem)
    total_mem=$(echo $memory_info | awk '{print $2}')
    used_mem=$(echo $memory_info | awk '{print $3}')
    memory_usage=$((used_mem * 100 / total_mem))
    
    log_info "Memory usage: ${memory_usage}% (${used_mem}MB/${total_mem}MB)"
    
    if [ $memory_usage -gt 90 ]; then
        log_error "Memory usage too high: ${memory_usage}%"
        return 1
    elif [ $memory_usage -gt 80 ]; then
        log_warn "Memory usage high: ${memory_usage}%"
    fi
    
    return 0
}

# 磁盘空间检查
disk_health_check() {
    log_info "Checking disk space..."
    
    # 检查应用数据目录
    if [ -d "/app/data" ]; then
        disk_usage=$(df /app/data | tail -1 | awk '{print $5}' | sed 's/%//')
        log_info "Disk usage: ${disk_usage}%"
        
        if [ $disk_usage -gt 90 ]; then
            log_error "Disk usage too high: ${disk_usage}%"
            return 1
        elif [ $disk_usage -gt 80 ]; then
            log_warn "Disk usage high: ${disk_usage}%"
        fi
    else
        log_warn "Data directory not found"
    fi
    
    return 0
}

# 进程检查
process_health_check() {
    log_info "Checking application process..."
    
    # 检查主进程是否运行
    if pgrep -f "uvicorn.*app.main:app" > /dev/null; then
        log_info "Application process is running"
        return 0
    else
        log_error "Application process not found"
        return 1
    fi
}

# 日志错误检查
log_health_check() {
    log_info "Checking for recent errors in logs..."
    
    if [ -d "/app/logs" ]; then
        # 检查最近5分钟内的错误日志
        error_count=$(find /app/logs -name "*.log" -newer /app/logs/.health_check_marker 2>/dev/null | wc -l)
        
        if [ $error_count -gt 0 ]; then
            log_warn "Found ${error_count} recent log files"
            
            # 检查错误级别日志
            error_lines=$(find /app/logs -name "*.log" -newer /app/logs/.health_check_marker -exec grep -l "ERROR\|CRITICAL" {} \; 2>/dev/null | wc -l)
            if [ $error_lines -gt 0 ]; then
                log_error "Found ${error_lines} log files with errors"
                return 1
            fi
        fi
    fi
    
    return 0
}

# 依赖服务检查
dependencies_health_check() {
    log_info "Checking external dependencies..."
    
    # 检查必要的Python包
    python3 -c "
import sys
required_packages = ['fastapi', 'pyttsx3', 'pydub', 'sqlalchemy']
missing_packages = []

for package in required_packages:
    try:
        __import__(package)
    except ImportError:
        missing_packages.append(package)

if missing_packages:
    print(f'Missing packages: {missing_packages}', file=sys.stderr)
    sys.exit(1)
else:
    print('All required packages available')
" 2>/dev/null
    
    if [ $? -eq 0 ]; then
        log_info "Dependencies check passed"
        return 0
    else
        log_error "Dependencies check failed"
        return 1
    fi
}

# 创建健康检查标记
create_marker() {
    touch /app/logs/.health_check_marker 2>/dev/null || true
}

# 清理函数
cleanup() {
    log_info "Cleaning up..."
}

# 主函数
main() {
    log_info "Starting comprehensive health check..."
    
    # 设置清理
    trap cleanup EXIT
    
    # 创建标记
    create_marker
    
    # 执行各项检查
    checks_passed=0
    total_checks=0
    
    # 基础检查
    log_info "=== Basic Health Checks ==="
    
    process_health_check
    if [ $? -eq 0 ]; then
        checks_passed=$((checks_passed + 1))
    fi
    total_checks=$((total_checks + 1))
    
    http_health_check
    if [ $? -eq 0 ]; then
        checks_passed=$((checks_passed + 1))
    fi
    total_checks=$((total_checks + 1))
    
    # 依赖检查
    log_info "=== Dependency Checks ==="
    
    dependencies_health_check
    if [ $? -eq 0 ]; then
        checks_passed=$((checks_passed + 1))
    fi
    total_checks=$((total_checks + 1))
    
    # 数据存储检查
    log_info "=== Data Storage Checks ==="
    
    database_health_check
    if [ $? -eq 0 ]; then
        checks_passed=$((checks_passed + 1))
    fi
    total_checks=$((total_checks + 1))
    
    redis_health_check
    if [ $? -eq 0 ]; then
        checks_passed=$((checks_passed + 1))
    fi
    total_checks=$((total_checks + 1))
    
    # 系统资源检查
    log_info "=== System Resource Checks ==="
    
    memory_health_check
    if [ $? -eq 0 ]; then
        checks_passed=$((checks_passed + 1))
    fi
    total_checks=$((total_checks + 1))
    
    disk_health_check
    if [ $? -eq 0 ]; then
        checks_passed=$((checks_passed + 1))
    fi
    total_checks=$((total_checks + 1))
    
    # 日志检查
    log_info "=== Log Health Checks ==="
    
    log_health_check
    if [ $? -eq 0 ]; then
        checks_passed=$((checks_passed + 1))
    fi
    total_checks=$((total_checks + 1))
    
    # 总结
    log_info "=== Health Check Summary ==="
    log_info "Passed: $checks_passed/$total_checks checks"
    
    if [ $checks_passed -eq $total_checks ]; then
        log_info "All health checks passed! ✓"
        exit 0
    else
        log_error "Some health checks failed! ✗"
        exit 1
    fi
}

# 检查命令
check_command curl
check_command python3
check_command pgrep
check_command free
check_command df

# 运行主函数
main "$@"
